# Author Zenos
# Create 2025/3/10 下午1:48
import numpy as np


def calculate_pixel_accuracy(matrix):
    """
        计算混淆矩阵的 Pixel Accuracy。
        参数：
        conf_matrix (ndarray): 混淆矩阵，形状为 [num_classes, num_classes]。
        返回：
        float: Pixel Accuracy，像素级别的准确率。
    """
    correct_pixels = np.trace(matrix)  # 正确预测的像素数
    total_pixels = matrix.sum()  # 总像素数
    return correct_pixels / total_pixels


def calculate_mean_pixel_accuracy(matrix):
    """
        计算混淆矩阵的平均像素准确率。
        参数：
        matrix (ndarray): 混淆矩阵，形状为 [num_classes, num_classes]。
        返回：
        float: 平均像素准确率，所有类别的准确率的平均值。
    """
    # 计算每个类别的正确预测像素数 (对角线)
    correct_pixels = np.diag(matrix)

    # 计算每个类别的总真实像素数 (每行的和)
    total_pixels_per_class = matrix.sum(axis=1)

    # 计算每个类别的像素准确率
    # 如果某类别没有真实样本，返回 NaN
    pixel_accuracy_per_class = np.where(total_pixels_per_class > 0, correct_pixels / total_pixels_per_class, np.nan)

    # 计算平均像素准确率，忽略 NaN 值
    average_pixel_accuracy = np.nanmean(pixel_accuracy_per_class)

    return average_pixel_accuracy, pixel_accuracy_per_class


def calculate_miou(conf_matrix):
    """
    计算均交并比 (Mean Intersection over Union, mIoU)，忽略背景类。
    参数:
    conf_matrix (ndarray): 混淆矩阵，形状为 [num_classes, num_classes]。
    返回:
    float: mIoU，所有类别 IoU 的均值，忽略背景类和没有真实样本的类别（返回 NaN）。
    """
    # 取对角线 (TP)
    intersection = np.diag(conf_matrix)

    # 计算分母 (TP + FP + FN)
    union = conf_matrix.sum(axis=1) + conf_matrix.sum(axis=0) - intersection

    # 创建一个掩码，排除背景类（假设类别0是背景）
    foreground_mask = np.arange(conf_matrix.shape[0]) != 0  # True for all non-background classes

    # 使用掩码排除背景类
    intersection = intersection[foreground_mask]
    union = union[foreground_mask]

    # 避免除零错误，如果该类别没有真实样本，则 IoU 设为 NaN
    iou_per_class = np.where(union > 0, intersection / union, np.nan)

    # 计算 mIoU（忽略 NaN 值）
    mean_iou = np.nanmean(iou_per_class)

    return mean_iou, iou_per_class


def foreground_frequency_weighted_intersection_over_union(confusion_matrix):
    """
    计算 Foreground Frequency Weighted IOU (FFW-IOU)

    参数：
    confusion_matrix: np.array，形状 (35, 35)，表示35个类别的混淆矩阵

    返回：
    FFW-IOU 值
    """
    num_classes = confusion_matrix.shape[0]

    # 计算每个类别的像素总数 t_i = sum(n_ij) (行和)
    t_i = confusion_matrix.sum(axis=1)  # shape (35,)

    # 计算每个类别的 IoU
    intersection = np.diag(confusion_matrix)  # 对角线上的值，正确分类的像素数
    union = t_i + confusion_matrix.sum(axis=0) - intersection  # t_i + sum(n_ji) - n_ii

    # 避免除零
    iou = np.where(union > 0, intersection / union, 0)

    # 只计算前景类别 (排除背景类别0)
    foreground_mask = np.arange(num_classes) != 0  # True for foreground classes

    # 计算 FFW-IOU
    foreground_pixel_count = t_i[foreground_mask].sum()  # ∑ t_k, k≠0
    if foreground_pixel_count == 0:
        return 0  # 如果没有前景像素，返回0

    ffw_iou = np.sum(t_i[foreground_mask] * iou[foreground_mask]) / foreground_pixel_count

    return ffw_iou


if __name__ == '__main__':
    foreground_frequency_weighted_intersection_over_union("sss")
